# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""数据流分析数据结构

提供 Def-Use Chain、逃逸分析等数据流分析的核心数据结构
"""

from dataclasses import dataclass, field
from enum import Enum
from typing import Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from mcpy.ir.base import IRNode
    from mcpy.ir.expressions.basic import IRName


class UseType(Enum):
    """变量使用类型"""

    CALL = "call"  # 函数调用: f(args)
    ASSIGNMENT = "assignment"  # 二次赋值: g = f
    RETURN = "return"  # 返回: return f
    CONTAINER = "container"  # 容器存储: [f] 或 funcs.append(f)
    PARAMETER = "parameter"  # 参数传递: func(f)
    ATTRIBUTE = "attribute"  # 属性访问: f.method()
    OTHER = "other"  # 其他


@dataclass
class Definition:
    """变量定义点

    Attributes:
        var_name: 变量名
        node: 定义节点（IRAssign, IRFor, IRFunctionDef 等）
        statement: 包含该定义的语句节点
        stmt_index: 语句在函数中的索引（用于比较先后顺序）
    """

    var_name: str
    node: object  # IRNode（避免循环导入）
    statement: object  # IRNode
    stmt_index: int


@dataclass
class Use:
    """变量使用点

    Attributes:
        var_name: 变量名
        node: 使用节点（通常是 IRName）
        statement: 包含该使用的语句节点
        stmt_index: 语句在函数中的索引
        use_type: 使用类型（调用、赋值、返回等）
    """

    var_name: str
    node: object  # IRName（避免循环导入）
    statement: object  # IRNode
    stmt_index: int
    use_type: UseType


@dataclass
class EscapeInfo:
    """变量逃逸分析信息

    Attributes:
        escapes: 是否逃逸
        reason: 逃逸原因（用于调试）
        escape_sites: 逃逸使用点列表
    """

    escapes: bool = False
    reason: str = ""
    escape_sites: list[Use] = field(default_factory=list)


@dataclass
class ClosureVariableUsage:
    """闭包变量的使用信息（数据流分析结果）

    Attributes:
        modified_after_capture: lambda 定义后是否被修改
        shared_by_multiple_lambdas: 是否被多个 lambda 共享
        outlives_lambda: 生命周期是否超过 lambda（预留）
        is_parameter: 是否是函数参数（参数不需要 shared_ptr）
    """

    modified_after_capture: bool = False
    shared_by_multiple_lambdas: bool = False
    outlives_lambda: bool = False
    is_parameter: bool = False


@dataclass
class DataFlowInfo:
    """函数级别的数据流信息

    为单个函数构建的数据流分析结果，包含：
    - 变量定义点和使用点
    - Def-Use Chain
    - 逃逸分析结果
    - 闭包变量使用信息

    Attributes:
        definitions: 变量定义点映射 {var_name: [Definition]}
        uses: 变量使用点映射 {var_name: [Use]}
        def_use_chain: Def-Use Chain {Definition: [Use]}
        escape_info: 逃逸分析结果 {var_name: EscapeInfo}
        closure_usage: 闭包变量使用信息 {var_name: ClosureVariableUsage}
    """

    definitions: dict[str, list[Definition]] = field(default_factory=dict)
    uses: dict[str, list[Use]] = field(default_factory=dict)
    def_use_chain: dict[Definition, list[Use]] = field(default_factory=dict)
    escape_info: dict[str, EscapeInfo] = field(default_factory=dict)
    closure_usage: dict[str, ClosureVariableUsage] = field(default_factory=dict)

    def add_definition(self, var_name: str, definition: Definition) -> None:
        """添加变量定义点"""
        if var_name not in self.definitions:
            self.definitions[var_name] = []
        self.definitions[var_name].append(definition)

    def add_use(self, var_name: str, use: Use) -> None:
        """添加变量使用点"""
        if var_name not in self.uses:
            self.uses[var_name] = []
        self.uses[var_name].append(use)

    def get_uses_after_def(self, var_name: str, def_stmt_index: int) -> list[Use]:
        """获取定义点之后的所有使用

        Args:
            var_name: 变量名
            def_stmt_index: 定义语句的索引

        Returns:
            定义之后的使用点列表
        """
        all_uses = self.uses.get(var_name, [])
        return [use for use in all_uses if use.stmt_index > def_stmt_index]

    def variable_escapes(self, var_name: str) -> bool:
        """检查变量是否逃逸

        Args:
            var_name: 变量名

        Returns:
            True 表示逃逸（需要堆分配或特殊处理）
        """
        escape_info = self.escape_info.get(var_name)
        return escape_info.escapes if escape_info else False

    def get_escape_reason(self, var_name: str) -> str:
        """获取变量逃逸原因（用于调试）

        Args:
            var_name: 变量名

        Returns:
            逃逸原因描述
        """
        escape_info = self.escape_info.get(var_name)
        return escape_info.reason if escape_info else "未逃逸"
